'''
Created on 2013. 8. 17.

@author: Hyunmin Kim, Euncheon Lim

JC69 model
JC69 is the simplest substitution model.
There are several assumptions.
It assumes equal base frequencies \left(\pi_A = \pi_G = \pi_C = \pi_T = {1\over4}\right) and equal mutation rates.
The only parameter of this model is therefore \mu, the overall substitution rate.
As previously mentioned, this variable becomes a constant when we normalize to the mean-rate to 1

'''
import numpy 
from scipy import linalg
import matplotlib.pyplot as plt
import string

class JC69Model:
    c_table = string.maketrans('ACGTacgt', '\0\1\2\3\0\1\2\3')
    def __init__(self):
        self.pi = 0.25 # an equilibrium vector of base (or character) frequencies
        self.t = numpy.linspace(0, 2, 100) # branch lengths
        
    ## log likelihood calculation
    ## exp(Q * t) =  P
    def __log_liklihood(self, t):
        P = linalg.expm(self.Q * t)
        sum_local = 0
        ## log (P(G) * P(G=>G) * P(G) * P(G=>A) ... P(C) * P(C=>C))
        for i in range(len(self.gene1)):
            sum_local += numpy.log(self.pi) + numpy.log(P[self.gene1[i], self.gene2[i]])
        return sum_local
    
    ## set the first gene sequence
    def set_gene1(self, s):
        self.gene1 = []
        s = string.translate(s, self.c_table)
        for c in s:
            self.gene1.append(ord(c))
    
    ## set the second gene sequence
    def set_gene2(self, s):
        self.gene2 = []
        s = string.translate(s, self.c_table)
        for c in s:
            self.gene2.append(ord(c))
    
    ## set the substitution_rate(mu)
    def set_substitution_rate(self, rate):
        self.mu = rate
        self.calculate_Q()
    
    ## calculate Q depending on mu
    def calculate_Q(self):
        self.Q = numpy.mat('[-3 1 1 1; 1 -3 1 1; 1 1 -3 1; 1 1 1 -3]')/4.0 * self.mu
        beta = 1/(-self.pi * self.Q.diagonal().sum())
        # normalize the Q independent of mu rates
        self.Q *= beta
        
    ## calculate the likelihood for each T point
    def calculate_likelihood_on_branch_length(self):
        self.y = [ self.__log_liklihood(x) for x in self.t ]
    
    def calculate_maximum_likelihood_point(self):
    	# Find the maximum y value
        max_y = max(self.y)
        max_t_index = self.y.index(max_y)
        # Find the x value corresponding to the maximum y value
        max_t = self.t[max_t_index]
        return max_t_index, max_t, max_y
        
    def plot_likelihood(self):
        ## find the maximum likelihood of branch length
        max_t_index, max_t, max_y= self.calculate_maximum_likelihood_point()
        
        plt.title('Log Maximally Likelihood Estimation')
        plt.xlabel('T(Branch Length)')
        plt.ylabel('P(Likelihood)')
        s = 'T={0}, P={1}'.format(*[max_t, max_y])
        plt.text(max_t, max_y + 0.1, s)
        plt.plot(self.t, self.y)
        plt.plot(self.t[max_t_index], self.y[max_t_index], 'rD')
        plt.show()

if  __name__ =='__main__':
    a_model = JC69Model()
    a_model.set_gene1('GGAC')
    a_model.set_gene2('GACC')
    a_model.set_substitution_rate(1)
    a_model.calculate_likelihood_on_branch_length()
    a_model.plot_likelihood()